;;; guile-openai --- An OpenAI API client for Guile
;;; Copyright © 2023 Andrew Whatson <whatson@tailcall.au>
;;;
;;; This file is part of guile-openai.
;;;
;;; guile-openai is free software: you can redistribute it and/or modify
;;; it under the terms of the GNU Affero General Public License as
;;; published by the Free Software Foundation, either version 3 of the
;;; License, or (at your option) any later version.
;;;
;;; guile-openai is distributed in the hope that it will be useful, but
;;; WITHOUT ANY WARRANTY; without even the implied warranty of
;;; MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
;;; Affero General Public License for more details.
;;;
;;; You should have received a copy of the GNU Affero General Public
;;; License along with guile-openai.  If not, see
;;; <https://www.gnu.org/licenses/>.

(define-module (openai utils event-stream)
  #:use-module (ice-9 rdelim)
  #:use-module (srfi srfi-41)
  #:export (port->line-stream
            line-stream->event-stream))

(define (read-crlf-line port)
  "Read a single line from PORT, delimited by either CR, LF, or CRLF.
The delimiter is included in the returned string.  Returns #f at
end-of-file."
  (let ((line (read-delimited "\r\n" port 'peek)))
    (and (string? line)
         (let* ((del1 (read-char port))
                (del1 (and (char? del1) del1))
                (del2 (and (eqv? del1 #\return)
                           (eqv? (peek-char port) #\newline)
                           (read-char port))))
           (cond ((and del1 del2)
                  (string-append line (string del1 del2)))
                 (del1
                  (string-append line (string del1)))
                 (else line))))))

(define (chomp-prefix prefix line)
  "Return the substring between PREFIX at the start of LINE, and a
CR/LF/CRLF delimiter at the end of LINE.  Returns #f if LINE doesn't
start with PREFIX."
  (and (string-prefix? prefix line)
       (let* ((len (string-length line))
              (del1 (string-ref line (- len 2)))
              (del2 (string-ref line (- len 1)))
              (offset (cond ((eqv? del1 #\return) 2)
                            ((eqv? del2 #\return) 1)
                            ((eqv? del2 #\newline) 1)
                            (else 0))))
         (substring/shared line
                           (string-length prefix)
                           (- len offset)))))

(define (port->line-stream port)
  "Return a stream which will yield each CR/LF/CRLF delimited line read
from PORT."
  (stream-let loop ()
    (let ((line (read-crlf-line port)))
      (if line
          (stream-cons line (loop))
          stream-null))))

(define (line-stream->event-stream strm)
  "Return a stream yielding the text/event-stream data payloads from the
line-stream STRM."
  (stream-let loop ((strm strm))
    (if (stream-null? strm)
        stream-null
        (let* ((line (stream-car strm))
               (data (chomp-prefix "data: " line))
               (rest (stream-cdr strm)))
          (cond ((not data)
                 (loop rest))
                ((string=? data "[DONE]")
                 stream-null)
                (else
                 (stream-cons data (loop rest))))))))
